/* 
 * Copyright 2014 Josef Hertl.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package caetus.pkcs11.key;

import caetus.pkcs11.template.Template;
import sun.security.pkcs11.wrapper.CK_MECHANISM;
import sun.security.pkcs11.wrapper.PKCS11;
import sun.security.pkcs11.wrapper.PKCS11Constants;
import sun.security.pkcs11.wrapper.PKCS11Exception;

/**
 *
 * @author Josef Hertl
 */
public class PrivateKey extends Key {

    private PublicKey publicKey;

    public PrivateKey(long sessionHandle, PKCS11 pkcs11, CipherEnum cipherEnum) {
        super(sessionHandle, pkcs11, cipherEnum);
    }

    @Override
    public byte[] encrypt(byte[] dataToEncrypt) throws PKCS11Exception {
        throw new UnsupportedOperationException("Private key cannot be used for encryption!");
    }

    @Override
    public byte[] decrypt(byte[] dataToDecrypt) throws PKCS11Exception {
        if (!super.isInitialized()) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
//            throw new IllegalStateException("The key is not initialized. Use generate() or create() first.");
        } else {
            CK_MECHANISM ck_mechanism = new CK_MECHANISM(cipherEnum.getCipherMechanism());
            pkcs11.C_DecryptInit(sessionHandle, ck_mechanism, keyHandle);
            byte[] decrypted = new byte[64];
            pkcs11.C_Decrypt(sessionHandle, dataToDecrypt, 0, dataToDecrypt.length, decrypted, 0, decrypted.length);
            return decrypted;
        }
    }

    @Override
    public byte[] wrapKey(Key keyToWrap) throws PKCS11Exception {
        throw new UnsupportedOperationException("Private key cannot be used for wrapping other keys!");
    }

    @Override
    public SecretKey unwrapSecretKey(byte[] toUnwrap, Template template, CipherEnum cipher) throws PKCS11Exception {
        if (!super.isInitialized()) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
//            throw new IllegalStateException("The key is not initialized. Use generate() or create() first.");
        } else {
            CK_MECHANISM unwrappingMechanism = new CK_MECHANISM(cipherEnum.getCipherMechanism());
            SecretKey unwrappedKey = new SecretKey(sessionHandle, pkcs11, cipher);
            unwrappedKey.setKeyHandleAndInitialize(pkcs11.C_UnwrapKey(sessionHandle, unwrappingMechanism, keyHandle, toUnwrap, template.finishAndGetAttributes()));
            return unwrappedKey;
        }
    }

    @Override
    public PrivateKey unwrapPrivateKey(byte[] toUnwrap, Template template, PublicKey correspondingPublicKey) throws PKCS11Exception {
        if (!super.isInitialized()) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
//            throw new IllegalStateException("The key is not initialized. Use generate() or create() first.");
        } else {
            CK_MECHANISM unwrappingMechanism = new CK_MECHANISM(cipherEnum.getCipherMechanism());
            PrivateKey unwrappedKey = new PrivateKey(sessionHandle, pkcs11, CipherEnum.RSA); //only asymmetric support is for RSA
            unwrappedKey.setKeyHandleAndInitialize(pkcs11.C_UnwrapKey(sessionHandle, unwrappingMechanism, keyHandle, toUnwrap, template.finishAndGetAttributes()));
            unwrappedKey.setPublicKey(correspondingPublicKey);
            return unwrappedKey;
        }
    }

    public PublicKey getPublicKey() {
        return publicKey;
    }

    public void setPublicKey(PublicKey publicKey) {
        this.publicKey = publicKey;
    }

    @Override
    public Key copyObject(Template template) throws PKCS11Exception {
        PrivateKey copiedKey = new PrivateKey(sessionHandle, pkcs11, cipherEnum);
        copiedKey.setPublicKey(publicKey);
        copiedKey.setKeyHandleAndInitialize(pkcs11.C_CopyObject(sessionHandle, keyHandle, template.finishAndGetAttributes()));
        return copiedKey;
    }

    @Override
    public Key deriveKey(CK_MECHANISM mechanism, Template template) throws PKCS11Exception {
        PrivateKey derivedKey = new PrivateKey(sessionHandle, pkcs11, cipherEnum);
        derivedKey.setPublicKey(publicKey);//this actually will not work, but should be good enough for the tests
        derivedKey.setKeyHandleAndInitialize(pkcs11.C_DeriveKey(sessionHandle, mechanism, keyHandle, template.finishAndGetAttributes()));
        return derivedKey;
    }
}
